import java.util.Arrays;
import java.util.Random;

public class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0) return new int[0];
        int p = selectK(arr, 0, arr.length - 1, k - 1, new Random());
        return Arrays.copyOf(arr, k);
    }

    /**
     * 求索引为 k 的元素
     *
     * @param arr
     * @param k
     * @return
     */
    public int selectK(int[] arr, int k) {
        return selectK(arr, 0, arr.length - 1, arr.length - k, new Random());
    }

    /**
     * 在arr[l,r]的范围里求解整个数组的第k小元素并返回
     * k是索引，即从0开始计算
     *
     * @param arr
     * @param l
     * @param r
     * @param k
     * @param random
     * @return
     */
    public int selectK(int[] arr, int l, int r, int k, Random random) {
        int p = partition(arr, l, r, random);
        if (k == p) return arr[p];
        if (k > p) return selectK(arr, p + 1, r, k, random);
        return selectK(arr, l, p - 1, k, random);
    }


    public int partition(int[] arr, int l, int r, Random random) {
        //生成 [l,r]之间的随机索引
        int p = random.nextInt(r - l + 1) + l;
        swap(arr, l, p);

        //arr[l+1...i-1]<=v   ; arr[j+1...r]>=v
        int i = l + 1, j = r;
        while (true) {
            while (i <= j && arr[i] < arr[l]) i++;
            while (i <= j && arr[j] > arr[l]) j--;
            if (i >= j) break;
            swap(arr, i, j);
            i++;
            j--;
        }
        swap(arr, l, j);
        return j;
    }

    private void swap(int[] arr, int i, int j) {
        int temp = arr[i];
        arr[i] = arr[j];
        arr[j] = temp;
    }

    public static void main(String[] args) {
        int[] arr = {3, 2, 3, 1, 2, 4, 5, 5, 6};
        Solution solution = new Solution();
        int i = solution.selectK(arr, 4);
        System.out.println("i = " + i);

        int[] arr2 = {0, 0, 2, 3, 2, 1, 1, 2, 0, 4};
        int[] leastNumbers = solution.getLeastNumbers(arr2, 10);
        for (int leastNumber : leastNumbers) {
            System.out.print(leastNumber + ",");
        }
    }
}
